
# 导入包
import os
import io
import json
import numpy as np
import pandas as pd
from tqdm import tqdm
import pandas as pd
tqdm.pandas(desc='pandas bar')
import glob
from pycococreatortools import pycococreatortools
from PIL import Image
import base64
from pathlib import Path

def img_tobyte(img_pil):
    '''
    该函数用于将图像转化为base64字符类型
    :param img_pil: Image类型
    :return base64_string: 字符串
    '''
    ENCODING = 'utf-8'
    img_byte = io.BytesIO()
    img_pil.save(img_byte, format='PNG')
    binary_str2 = img_byte.getvalue()
    imageData = base64.b64encode(binary_str2)
    base64_string = imageData.decode(ENCODING)
    return base64_string

def mask_to_json(mask_path, img_root, dst_root, code):
    class_names = ['_background_', 'defect'] # 分别表示label标注图中1对应basketball，2对应person。
    if code is None:
        code = Path(mask_path).parent.name

    Json_output = {
        "version": "5.0.1",
        "flags": {},
        "fillColor": [255, 0, 0, 128],
        "lineColor": [0, 255, 0, 128],
        "imagePath": {},
        "shapes": [],
        "imageData": {}}
    # print(mask_path)

    img_path = os.path.join(img_root, code, Path(mask_path).with_suffix('.jpg').name)
    Json_output["imagePath"] = img_path
    # 打开原图并将其转化为labelme json格式
    # image = Image.open(img_path)
    # imageData = img_tobyte(image)
    # Json_output["imageData"] = imageData
    # 获得注释的掩码
    binary_mask = np.asarray(np.array(Image.open(mask_path))).astype(np.uint8) // 255
    # 分别对掩码中的label结果绘制边界点
    try:
        for i in np.unique(binary_mask):
            if i != 0:
                temp_mask = np.where(binary_mask == i, 1, 0)
                segmentation = pycococreatortools.binary_mask_to_polygon(temp_mask, tolerance=1) # tolerancec参数控制无误差
                for item in segmentation:
                    # if (len(item) > 10):
                    list1 = []
                    for j in range(0, len(item), 2):
                        list1.append([item[j], item[j + 1]])
                    label = class_names[i]  #
                    seg_info = {'points': list1, "fill_color": None, "line_color": None, "label": label,
                                "shape_type": "polygon", "flags": {}}
                    Json_output["shapes"].append(seg_info)
        Json_output["imageHeight"] = binary_mask.shape[0]
        Json_output["imageWidth"] = binary_mask.shape[1]
        # 保存在根目录下的json文件中
        json_path = os.path.join(dst_root, code, Path(mask_path).with_suffix('.json').name)
        os.makedirs(Path(json_path).parent, exist_ok=True)
        with open(json_path, 'w') as output_json_file:
            json.dump(Json_output, output_json_file, indent=4)
    except:
        return 

def main():
    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False
    
    img_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/cropped_img_0409'
    mask_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/cropped_mask_0409'
    dst_root = img_root  # r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6006/dst'

    df = pd.DataFrame()
    df['mask_path'] = glob.glob(os.path.join(mask_root, "*/*.png"))
    df['code'] = df['mask_path'].progress_apply(lambda x: Path(x).parent.name)
    df.parallel_apply(lambda x: mask_to_json(mask_path=x['mask_path'], img_root=img_root, dst_root=dst_root, code=x['code']) , axis=1)

    print()

if __name__=='__main__':
    main()
